﻿#include "FocusMeasure.h"

static double roberts(cv::Mat &grayImage)
{
    double result = 0;

    for (int j = 0; j < grayImage.rows - 1; j++)
    { // for all rows
        const uchar* current = grayImage.ptr<const uchar>(j); // current row
        const uchar* next = grayImage.ptr<const uchar>(j + 1); // next row
        for (int i = 0; i < grayImage.cols - 1; i++)
        {
            // 1   0
            // 0  -1
            double Gx = current[i] - next[i + 1];
            // 0  -1
            // 1   0
            double Gy = next[i] - current[i + 1];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return result/(grayImage.rows * grayImage.cols);
}

static double fastRoberts(cv::Mat &grayImage)
{
    double result = 0;

    for (int j = 0; j < grayImage.rows - 1; j += 2)
    { // for all rows
        const uchar* current = grayImage.ptr<const uchar>(j); // current row
        const uchar* next = grayImage.ptr<const uchar>(j + 1); // next row
        for (int i = 0; i < grayImage.cols - 1; i += 2)
        {
            // 1   0
            // 0  -1
            double Gx = current[i] - next[i + 1];
            // 0  -1
            // 1   0
            double Gy = next[i] - current[i+1];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return 4*result/(grayImage.rows * grayImage.cols);
}

static double sobel(cv::Mat &grayImage)
{
    double result = 0;
    for (int j = 1; j < grayImage.rows - 1; j++)
    { // for all rows
        const uchar* previous = grayImage.ptr<const uchar>(j-1); // previous row
        const uchar* current = grayImage.ptr<const uchar>(j); // current row
        const uchar* next = grayImage.ptr<const uchar>(j+1); // next row

        for (int i = 1; i < grayImage.cols - 1; i++)
        {
            // -1  0  1
            // -2  0  2
            // -1  0  1
            double Gx = previous[i + 1] + next[i + 1]
                    - previous[i - 1] - next[i - 1]
                    + 2 * current[i + 1] - 2 * current[i - 1];

            // -1 -2 -1
            //  0  0  0
            //  1  2  1
            double Gy = next[i - 1] + next[i + 1]
                    - previous[i - 1] - previous[i + 1]
                    + 2 * previous[i] - 2 * next[i];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return result/(grayImage.rows * grayImage.cols);
}


static double fastSobel(cv::Mat &grayImage)
{
    double result = 0;
    for (int j = 1; j < grayImage.rows - 1; j += 3)
    { // for all rows
        const uchar* previous = grayImage.ptr<const uchar>(j-1); // previous row
        const uchar* current = grayImage.ptr<const uchar>(j); // current row
        const uchar* next = grayImage.ptr<const uchar>(j+1); // next row

        for (int i = 1; i < grayImage.cols - 1; i += 3)
        {
            // -1  0  1
            // -2  0  2
            // -1  0  1
            double Gx = previous[i + 1] + next[i + 1]
                    - previous[i - 1] - next[i - 1]
                    + 2 * current[i + 1] - 2 * current[i - 1];

            // -1 -2 -1
            //  0  0  0
            //  1  2  1
            double Gy = next[i - 1] + next[i + 1]
                    - previous[i - 1] - previous[i + 1]
                    + 2 * previous[i] - 2 * next[i];
            result += Gx * Gx + Gy + Gy;
        }
    }
    return 9*result/(grayImage.rows * grayImage.cols);
}

static double laplace(cv::Mat &grayImage)
{
    double result = 0.0;
    for (int j = 1; j < grayImage.rows - 1; j++)
    { // for all rows
        const uchar* previous = grayImage.ptr<const uchar>(j-1); // previous row
        const uchar* current = grayImage.ptr<const uchar>(j); // current row
        const uchar* next = grayImage.ptr<const uchar>(j+1); // next row
        for (int i = 1; i < grayImage.cols - 1; i++)
        {
            // 0  -1   0
            //-1   4  -1
            // 0  -1   0
            result +=  4 * current[i] - current[i - 1] - current[i + 1]
                    - previous[i] - next[i];
        }
    }
    return result/(grayImage.rows * grayImage.cols);
}

static double fastLaplace(cv::Mat &grayImage)
{
    double result = 0.0;
    for (int j = 1; j < grayImage.rows - 1; j += 3)
    { // for all rows
        const uchar* previous = grayImage.ptr<const uchar>(j-1); // previous row
        const uchar* current = grayImage.ptr<const uchar>(j); // current row
        const uchar* next = grayImage.ptr<const uchar>(j+1); // next row
        for (int i = 1; i < grayImage.cols - 1; i += 3)
        {
            // 0  -1   0
            //-1   4  -1
            // 0  -1   0
            result +=  4 * current[i] - current[i - 1] - current[i + 1]
                    - previous[i] - next[i];
        }
    }
    return 9*result/(grayImage.rows * grayImage.cols);
}

static double discreteVariance(cv::Mat &grayImage)
{
    double mean = 0.0;
    double square = 0.0;
    double result = 0.0;
    for (int j = 0; j < grayImage.rows - 1; j ++)
    { // for all rows
        const uchar* current = grayImage.ptr<const uchar>(j); // current row
        double mean_line = 0.0;
        double square_line = 0.0;
        for (int i = 0; i < grayImage.cols - 1; i ++)
        {
            mean_line += current[j];
            square_line += current[j] * current[j];
        }
        mean_line /= grayImage.cols;
        square_line /= grayImage.cols;
        mean += mean_line;
        square += square_line;

    }
    mean /= grayImage.rows;
    square /= grayImage.rows;
    result = square - mean * mean;
    return result;
}

FocusMeasure::FocusMeasure()
{
    setMeasureMethod(ROBERTS);
    setZoomFactor(0);
}

FocusMeasure::~FocusMeasure()
{

}

void FocusMeasure::setMeasureMethod(MeasureMethod method)
{
    switch(method)
    {
    case ROBERTS:
        m_measureFunc = roberts;
        break;
    case FAST_ROBERTS:
        m_measureFunc = fastRoberts;
        break;
    case SOBEL:
        m_measureFunc = sobel;
        break;
    case FAST_SOBEL:
        m_measureFunc = fastSobel;
        break;
    case LAPLACE:
        m_measureFunc = laplace;
        break;
    case FAST_LAPLACE:
        m_measureFunc = fastLaplace;
        break;
    case DISCRETE_VARIANCE:
        m_measureFunc = discreteVariance;
        break;
    default:
        m_measureFunc = roberts;
    }
}

cv::Mat FocusMeasure::preprocess(QImage &image)
{
    cv::Mat mat = QImage2cvMat(image);
    cv::Mat grayImage;
    cv::cvtColor(mat, grayImage, cv::COLOR_BGR2GRAY);
    if(m_zoomFactor > 1)
    {
        cv::resize(grayImage, grayImage, cv::Size(image.width() / m_zoomFactor, image.height() / m_zoomFactor));
    }
    return grayImage;
}
